#!/usr/bin/env python3
import pysam
import pandas as pd


def reverse_seq(seq):
    COMPLEMENT = {
        "A": "T", "C": "G", "G": "C", "T": "A",
        "N": "N", "W": "W", "S": "S", "K": "M",
        "M": "K", "Y": "R", "R": "Y", "H": "D",
        "B": "V", "V": "B", "D": "H"
    }

    return "".join(reversed([COMPLEMENT[i] for i in seq]))


def is_overlap(interval1, interval2):
    x1, y1 = interval1
    x2, y2 = interval2
    if y1 <= x2 or y2 <= x1:
        return False
    else:
        return True
    

def split_discontinuous(discontinuous_list):
    continuous_list = []
    for p in discontinuous_list:
        if not continuous_list:
            continuous_list.append(p)
        elif p - continuous_list[-1] == 1:
            continuous_list.append(p)
        else:
            yield continuous_list
            continuous_list = [p]
    if continuous_list:
        yield continuous_list


def table2df(table, header=0, sep=None):
    if table.endswith(".csv"):
        deli = sep if sep else ","
        df = pd.read_csv(table, delimiter=deli, header=header)
    elif table.endswith(".xlsx"):
        df = pd.read_excel(table, header=header)
    else:
        deli = sep if sep else "\t"
        df = pd.read_csv(table, delimiter=deli, header=header)
    
    return df


def df2table(table, out_df, index=False, header=True, sep="\t"):
    if table.endswith(".csv"):
        out_df.to_csv(table, index=index, header=header, sep=",")
    elif table.endswith(".xlsx"):
        out_df.to_excel(table, index=index, header=header)
    else:
        out_df.to_csv(table, index=index, header=header, sep=sep)
        

class GenomeFetcher:
    def __init__(self, genome):
        self.genome = pysam.FastaFile(genome)

    def fetch(self, chrom, blocks, strand):
        sequence = ""
        chromLen = self.get_chrom_length(chrom)
        for block in blocks:
            start = int(block[0]) if int(block[0]) > 0 else 0
            end = int(block[1]) if int(block[1]) < chromLen else chromLen
            if end <= start:
                continue
            sequence += self.genome.fetch(chrom, start, end)
        transcript = sequence if strand == '+' else reverse_seq(sequence)

        return transcript

    def get_chrom_length(self, chrom):
        return self.genome.get_reference_length(chrom)
